import os

from PIL import Image

from option import parse_args


def get_min_width(video_path):
    w_min, w_max = 300, 0
    for image_name in os.listdir(video_path):
        with Image.open(os.path.join(video_path, image_name)) as img:
            w_min = min(w_min, img.size[0])
            w_max = max(w_max, img.size[0])
    return w_min, w_max


def get_width_less_than_100(video_path):
    width = {}
    for image_name in os.listdir(video_path):
        with Image.open(os.path.join(video_path, image_name)) as img:
            if img.size[0] < 100:
                if img.size[0] not in width:
                    width[img.size[0]] = 1
                else:
                    width[img.size[0]] += 1
    return width


def get_video_name_set_by_mode(args, mode):
    video_name_set = set()
    video_names = open(os.path.join(args.root_dir, '{}.txt'.format(mode)), 'r').readlines()
    for video_name in video_names:
        video_name_split = video_name.split(',')
        video_name_set.add((video_name_split[0], video_name_split[1], video_name_split[2].strip()))
    return video_name_set


def get_video_name_set(args):
    video_name_set = set()
    video_name_set.update(get_video_name_set_by_mode(args, 'train'))
    video_name_set.update(get_video_name_set_by_mode(args, 'test'))
    video_name_set.update(get_video_name_set_by_mode(args, 'val'))
    return video_name_set


def abnormal_crops(args, func, template, ):
    crops = os.path.join(args.root_dir, 'crops')
    video_name_set = get_video_name_set(args)
    for video_name_pair in video_name_set:
        video_path = os.path.join(crops, video_name_pair[1])
        # 1. Count with min and max
        w_min, w_max = get_min_width(video_path)

        if w_min < 100:
            func(template.format(video_name_pair[0], video_name_pair[1], video_name_pair[2], w_min, w_max))

        # 2. Count with dict
        # width = get_width_less_than_100(video_path)
        # if len(width) != 0:
        #     print('{},{}'.format(video_name, width))


def main():
    args = parse_args()
    assert os.path.exists(args.root_dir)
    args.save = False
    if args.save:
        with open('abnormal_width_statistics.dat', 'w') as f:
            abnormal_crops(args, f.write, '{},{},{},{},{}\n')
    else:
        abnormal_crops(args, print, '{},{},{},{},{}')


if __name__ == '__main__':
    main()
